package com.bonc.classifier.maxent;

import java.io.File;
import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.ObjectOutputStream;
import java.io.PrintStream;
import java.util.ArrayList;
import java.util.Collections;
import java.util.Comparator;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Set;

import com.bonc.utilities.EventUtility;
import com.bonc.vectorspace.model.EventCorpus;
import com.bonc.vectorspace.model.EventDocument;
import com.bonc.vectorspace.model.VectorSpaceModel;

import cc.mallet.classify.Classifier;
import cc.mallet.classify.ClassifierTrainer;
import cc.mallet.classify.MaxEnt;
import cc.mallet.classify.MaxEntTrainer;
import cc.mallet.types.Alphabet;
import cc.mallet.types.FeatureSelection;
import cc.mallet.types.FeatureVector;
import cc.mallet.types.InfoGain;
import cc.mallet.types.Instance;
import cc.mallet.types.InstanceList;
import cc.mallet.types.Label;
import cc.mallet.types.LabelAlphabet;
import cc.mallet.types.Labeling;
import cc.mallet.util.FileUtils;

public class MaxentTrainer{
	
	private Set<String> featureTerms;
	
	private Classifier classifier;
	
	public MaxentTrainer(){
		
	}
	
	public MaxentTrainer(String trainingDirname, String modelFileName){

		try {
			List<EventDocument> documents = EventUtility.convertFiles2Document(trainingDirname);
			EventCorpus trainingCorpus = new EventCorpus(documents);
			File modelFile = new File(modelFileName);
			classifier = this.TrainMaxent(trainingCorpus, modelFile);
		} catch (IOException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		}
	}
	
	/**
	 * Use information gain to rank and select features
	 * @param instances
	 * @return 
	 */
	public static void pruneByInfoGain(InstanceList instances){
		InfoGain ig = new InfoGain (instances);
		int numfeatures = (int) (instances.getDataAlphabet().size() * 0.9);
		FeatureSelection fs = new FeatureSelection (ig, numfeatures);
		for (int ii = 0; ii < instances.size(); ii++) {
			Instance instance = instances.get(ii);
			FeatureVector fv = (FeatureVector) instance.getData();
			FeatureVector fv2 = FeatureVector.newFeatureVector (fv, instances.getDataAlphabet(), fs);
			instance.unLock();
			instance.setData(fv2);
		}
	}
	
	
	// change this accordingly
	static double Gaussian_Variance = 1.0;
	
	// in the training feature table
	// Lines should be formatted as:                                                                   
    //                                                                                                 
    //   [name] [label] [data ... ]                                                                    
    //                                                                                                          
//	public Classifier TrainMaxent(File trainingFilename, File modelFile) throws IOException {
//		// build data input pipe
//		ArrayList<Pipe> pipes = new ArrayList<Pipe>();
//		
//		// define pipe
//		// the features in [data ...] should like: feature:value
//		pipes.add(new Target2Label());
//		pipes.add(new Csv2FeatureVector());
//		
//		Pipe pipe = new SerialPipes(pipes);
//		pipe.setTargetProcessing(true);
//		
//		// read data
//		InstanceList trainingInstances = new InstanceList(pipe);
//		FileReader training_file_reader = new FileReader(trainingFilename);
//		CsvIterator reader =
//	            new CsvIterator(training_file_reader,
//	                            "(\\w+)\\s+(\\S+)\\s+(.*)",
//	                            3, 2, 1);  // (data, label, name) field indices    
//		trainingInstances.addThruPipe(reader);
//		training_file_reader.close();
//		
//		// prune by info gain
//		// pruneByInfoGain(trainingInstances);
//		
//		// calculate running time
//		long startTime = System.currentTimeMillis();
//		PrintStream temp = System.err;
//		System.setErr(System.out);
//		
//		// train a Maxent classifier (could be other classifiers)
//		ClassifierTrainer<MaxEnt> trainer = new MaxEntTrainer(Gaussian_Variance);
//		Classifier classifier = trainer.train(trainingInstances);
//		
//		System.setErr(temp);
//		// calculate running time
//		long endTime = System.currentTimeMillis();
//		long totalTime = endTime - startTime;
//		System.out.println("Total training time: " + totalTime);
//		
//		// write model
//		ObjectOutputStream oos = new ObjectOutputStream(new FileOutputStream(modelFile));
//	    oos.writeObject(classifier);
//	    oos.close();
//	   
//	    
//	    return classifier;
//	}

	/**
	 * print the feature weights for the MaxEnt model
	 * @param maxEnt
	 * @param out
	 */
	public static void printFeatureWeights(MaxEnt maxEnt, File out){
		PrintStream weights_out;
		try
		{
			weights_out = new PrintStream(out);
			maxEnt.print(weights_out);
			weights_out.close();
			
			String[] lines = FileUtils.readFile(out);
			weights_out = new PrintStream(out); 
			Map<String, Double> map = null;
			String label = null;
			for(String line : lines)
			{
				if(line.startsWith("FEATURES"))
				{
					if(map != null)
					{
						// sort map
						List<Entry<String, Double>> entrylist = new ArrayList<Entry<String, Double>>();
						entrylist.addAll(map.entrySet());
						Collections.sort(entrylist, new Comparator<Entry<String, Double>>()
								{
//									@Override
									public int compare(
											Entry<String, Double> o1,
											Entry<String, Double> o2)
									{
										return -1 * o1.getValue().compareTo(o2.getValue());
									}
								}
								);
						// print out a feature table for a class
						for(Map.Entry<String, Double> entry : entrylist)
						{
							weights_out.println(label + "\t" + entry.getKey() + "\t" + entry.getValue());
						}
						weights_out.println();
					}
					// create a new map
					map = new HashMap<String, Double>();
					String[] cols = line.split("\\s+");
					label = cols[cols.length - 1];
				}
				else
				{
					if(map != null)
					{
						String[] cols = line.trim().split("\\s+");
						if(cols.length == 2)
						{
							String feat = cols[0];
							Double weight = Double.parseDouble(cols[1]);
							map.put(feat, weight);
						}
					}
				}
			}
			if(map != null)
			{
				// sort map
				List<Entry<String, Double>> entrylist = new ArrayList<Entry<String, Double>>();
				entrylist.addAll(map.entrySet());
				Collections.sort(entrylist, new Comparator<Entry<String, Double>>()
						{
//							@Override
							public int compare(
									Entry<String, Double> o1,
									Entry<String, Double> o2)
							{
								return -1 * o1.getValue().compareTo(o2.getValue());
							}
						}
						);
				// print out a feature table for a class
				for(Map.Entry<String, Double> entry : entrylist)
				{
					weights_out.println(label + "\t" + entry.getKey() + "\t" + entry.getValue());
				}
				weights_out.println();
			}
		} 
		catch (FileNotFoundException e)
		{
			e.printStackTrace();
		} 
		catch (IOException e)
		{
			e.printStackTrace();
		}
	}

    //predict & evaluate
    public static String predict(Classifier myClassifier,Instance testInstance){
        Labeling labeling = myClassifier.classify(testInstance).getLabeling();
        Label label = labeling.getBestLabel();
        return (String)label.getEntry();
    }
    
    public Labeling predict(Instance testInstance){
        Labeling labeling = classifier.classify(testInstance).getLabeling();
//        Label label = labeling.getBestLabel();
//        return (String)label.getEntry();
        return labeling;
    }
    

	public Classifier TrainMaxent(EventCorpus trainingCorpus, File modelFile) throws IOException {
		VectorSpaceModel vectorSpace = new VectorSpaceModel(trainingCorpus);
		
		this.setFeatureTerms(trainingCorpus.getInvertedIndex().keySet());

		//maxent
        Alphabet featureAlphabet = new Alphabet();//特征词典
        LabelAlphabet targetAlphabet = new LabelAlphabet();//类标词典
        targetAlphabet.lookupIndex("pos");
        targetAlphabet.lookupIndex("neg");
        targetAlphabet.stopGrowth();
        for (String term : featureTerms) {
        	featureAlphabet.lookupIndex(term);	       	
        }

        InstanceList allInstances = new InstanceList (featureAlphabet,targetAlphabet);//实例集对象
        
        int featuresize = featureTerms.size();

		for (EventDocument document : trainingCorpus.getDocuments()) {
//			System.out.print("document "+(document.getFileName()+"===="));
			HashMap<String, Double> weights = vectorSpace.getTfIdfWeights().get(document);
			double[] featureValues1 = new double[featuresize];
			int j = 0;
			for (String term : featureTerms) {				
				Double weight = weights.get(term);
				if(weight !=null ){
					featureValues1[j] = weight.doubleValue();
				}else{
					featureValues1[j] = 0.0;
				}
				j++;				
			}
			System.out.println();			
            
            FeatureVector featureVector = new FeatureVector(featureAlphabet,featureTerms.toArray(new String[featuresize]),featureValues1);//change list to array           
            System.out.println("test1=="+document.getTag());
            Instance instance = new Instance (featureVector,targetAlphabet.lookupLabel(document.getTag()), document.getDocId(),null);
//            Instance instance = new Instance (featureVector,targetAlphabet.lookupLabel(targetValue.get(i)), document.getFileName(),null);
            
            allInstances.add(instance);
		}

		// prune by info gain
		// pruneByInfoGain(trainingInstances);
		
		// calculate running time
		long startTime = System.currentTimeMillis();
		PrintStream temp = System.err;
		System.setErr(System.out);
		
		double Gaussian_Variance = 1.0;
		ClassifierTrainer<MaxEnt> trainer = new MaxEntTrainer(Gaussian_Variance);		
		Classifier classifier = trainer.train(allInstances); 
		
		System.setErr(temp);
		// calculate running time
		long endTime = System.currentTimeMillis();
		long totalTime = endTime - startTime;
		System.out.println("Total training time: " + totalTime);
		
		// write model
		ObjectOutputStream oos = new ObjectOutputStream(new FileOutputStream(modelFile));
	    oos.writeObject(classifier);
	    oos.close();
	    
		return classifier;		
	}

	public Set<String> getFeatureTerms() {
		return featureTerms;
	}

	public void setFeatureTerms(Set<String> featureTerms) {
		this.featureTerms = featureTerms;
	}

	public Classifier getClassifier() {
		return classifier;
	}

	public void setClassifier(Classifier classifier) {
		this.classifier = classifier;
	}	
	
//	public static void main (String[] args) throws Exception 
//	{
//		if(args.length < 2)
//		{
//			System.out.println("Usage:");
//			System.out.println("args[0] : feature table file");
//			System.out.println("args[1] : model file name");
//			System.exit(-1);
//		}		
//		
//		File featureTable = new File(args[0]);
//		File modelFile = new File(args[1]);
//		
//		// train the model
//		TrainMaxent(featureTable, modelFile); 
//	}
}
